# -*- coding: utf-8 -*-
import logging
import os
import json
import subprocess
import shutil
from boto3.session import Session
from botocore.exceptions import ClientError

LOGGER = logging.getLogger()

FFMPEG = "/opt/ffmpeg/ffmpeg"


def get_fileNameExt(filename):
    (_, tempfilename) = os.path.split(filename)
    (shortname, extension) = os.path.splitext(tempfilename)
    return shortname, extension


def handler(event, context):
    LOGGER.info(event)

    access_key = os.getenv("ext_cf_zos_ak")
    secret_key = os.getenv("ext_cf_zos_sk")
    endpoint = os.getenv("ext_cf_zos_endpoint")
    if access_key is None:
        return "env check failed: zos ak is None"
    if secret_key is None:
        return "env check failed: zos sk is None"
    if endpoint is None:
        return "env check failed: zos endpoint is None"

    session = Session(access_key, secret_key)
    s3_client = session.client("s3", endpoint_url=endpoint)

    evt = json.loads(event)
    bucket = evt["bucket"]
    remote_file = evt["object"]
    output_dir = evt["output_dir"]
    dst_format = evt['dst_format']
    shortname, ext = get_fileNameExt(remote_file)

    local_save_dir = "/tmp"
    local_save_file = local_save_dir + "/" + shortname + "." + ext

    # 下载源文件
    try:
        s3_client.download_file(bucket, remote_file, local_save_file)
    except ClientError as e:
        return "download source file from zos failed: %s" % e

    # m3u8 特殊处理
    rid = context.request_id
    if dst_format == "m3u8":
        return handle_m3u8(rid, s3_client, bucket, local_save_file, shortname, output_dir)
    else:
        return handle_common(rid, s3_client, bucket, local_save_file, shortname, output_dir, dst_format)


def handle_m3u8(request_id, s3_client, bucket, input_path, shortname, output_dir):
    ts_dir = '/tmp/ts'
    if os.path.exists(ts_dir):
        shutil.rmtree(ts_dir)
    os.mkdir(ts_dir)
    ts_file = os.path.join('/tmp', shortname + '.ts')
    split_transcoded_filepath = os.path.join(ts_dir, shortname + '_%03d.ts')
    cmd1 = [FFMPEG, '-y', '-i', input_path, '-c:v', 'libx264', ts_file]
    cmd2 = [FFMPEG, '-y', '-i', ts_file, '-c', 'copy', '-map',  '0',  '-f', 'segment',
            '-segment_list', os.path.join(ts_dir, 'playlist.m3u8'), '-segment_time', '10', split_transcoded_filepath]

    try:
        subprocess.run(
            cmd1, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)

        subprocess.run(
            cmd2, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)

        for filename in os.listdir(ts_dir):
            filepath = os.path.join(ts_dir, filename)
            filekey = os.path.join(output_dir, shortname, filename)
            s3_client.upload_file(filepath, bucket, filekey)
            os.remove(filepath)
            print("Uploaded {} to {}".format(filepath, filekey))

    except subprocess.CalledProcessError as exc:
        # if transcode fail，trigger invoke dest-fail function
        raise Exception(request_id +
                        " transcode failure, detail: " + str(exc))

    finally:
        if os.path.exists(ts_dir):
            shutil.rmtree(ts_dir)

        # remove ts 文件
        if os.path.exists(ts_file):
            os.remove(ts_file)

    return {"msg": "success"}


def handle_common(request_id, s3_client, bucket, input_path, shortname, output_dir, dst_format):
    transcoded_filepath = os.path.join('/tmp', shortname + '.' + dst_format)
    if os.path.exists(transcoded_filepath):
        os.remove(transcoded_filepath)
    cmd = [FFMPEG, "-y", "-i", input_path, transcoded_filepath]
    try:
        subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
        s3_client.upload_file(transcoded_filepath, bucket, os.path.join(output_dir, shortname + '.' + dst_format))
    except subprocess.CalledProcessError as exc:
        # if transcode fail，trigger invoke dest-fail function
        raise Exception(request_id +
                        " transcode failure, detail: " + str(exc))
    finally:
        if os.path.exists(transcoded_filepath):
            os.remove(transcoded_filepath)

    return {"msg": "success"}
